import numpy as np
import torch
import scipy.sparse.linalg as sla
from pathlib import Path
from .numerical_method import reaction_A, reaction_b, lap_A, lap_b 
from .delta import normal

def func_a(x, y):
    return 1 + 2 * y**2

def func_r(x, y):
    return 1 + x ** 2

def func_dady(x, y):
    return 4 * y

def func_dadx(x, y):
    return 0 * x

# def func_u(x, y):
#     return np.e ** (-(x**2 + 2 * y**2 + 1))

# def force(x, y):
#     return (-8*x**2*y**2 - 3*x**2 - 32*y**4 + 12*y**2 + 7) * func_u(x, y)

def gen_data(grid_num, data_path='./data', train_N=2500, val_N=100, a=1, label=''):
    h = 2 * a / (grid_num - 1)
    x = np.linspace(-a, a, grid_num)
    y = np.linspace(-a, a, grid_num)
    xx, yy = np.meshgrid(x, y)
    
    locs_x = np.random.uniform(-a, a, train_N)
    locs_y = np.random.uniform(-a, a, train_N)
    locs_train = np.stack((locs_x, locs_y), axis=1)

    locs_x = np.linspace(-a+2*h, a-2*h, val_N)
    locs_y = np.linspace(-a+2*h, a-2*h, val_N)
    locs_val = np.stack((locs_x, locs_y), axis=1)

    locs = np.vstack((locs_train, locs_val))
    # mkdir
    save_path = f'{data_path}/{grid_num}/{label}'
    p = Path(save_path)
    if not p.is_dir():  p.mkdir(parents=True)
    
    # generate data
    _gen_lap_real(save_path, xx, yy, h, locs, train_N)
    _gen_reaction_real(save_path, xx, yy, h, locs, train_N)
    _gen_r_input(save_path, xx, yy, locs, train_N)
    _gen_f_input(save_path, xx, yy, locs, train_N, normal, h)

    mat_a = func_a(xx, yy)
    np.save(f'{data_path}/{grid_num}/a.npy', mat_a)

    mat_r = func_r(xx, yy)
    np.save(f'{data_path}/{grid_num}/r.npy', mat_r)
    return True

def _gen_r_input(path, xx, yy, locs, N):
    data = np.zeros((len(locs), *xx.shape))
    for i, p in enumerate(locs):
        data[i] = _r(xx, yy, p)
    np.save(f'{path}/R.npy', data[:N])
    np.save(f'{path}/valR.npy', data[N:])
    return True

def _r(xx, yy, p):
    r = (np.sqrt((xx-p[0])**2 + (yy-p[1])**2))
    return (r.max() - r) / (r.max() - r.min())

def _gen_lap_real(path, xx, yy, h, locs, N):
    n = xx.shape[0]
    A = lap_A(n)
    # B = np.zeros((n**2, len(locs)))
    uu = []
    for i, p in enumerate(locs):
        source = normal(xx, yy, h, p)
        b = lap_b(source, h)
        uu.append(sla.spsolve(A, b).reshape((-1, n, n)))
        
    np.save(f'{path}/U.npy', np.stack(uu[:N], axis=0))
    np.save(f'{path}/valU.npy', np.stack(uu[N:], axis=0))
    
def _gen_reaction_real(path, xx, yy, h, locs, N):
    n = xx.shape[0]
    A = reaction_A(n)
    uu = []
    for i, p in enumerate(locs):
        source = normal(xx, yy, h, p)
        b = reaction_b(source, h)
        uu.append(sla.spsolve(A, b).reshape((-1, n, n)))
        
    np.save(f'{path}/rU.npy', np.stack(uu[:N], axis=0))
    np.save(f'{path}/valrU.npy', np.stack(uu[N:], axis=0))
    
def _gen_f_input(path, xx, yy, locs, N, func, h):
    data = np.zeros((len(locs), *xx.shape))
    for i, p in enumerate(locs):
        data[i] = func(xx, yy, h, p)
    np.save(f'{path}/F.npy', data[:N])
    np.save(f'{path}/valF.npy', data[N:])
    return True

def assemble_F(n, points):
    h = 2 / (n - 1)
    x = np.linspace(-1, 1, n)
    y = np.linspace(-1, 1, n)
    xx, yy = np.meshgrid(x, y)
    dataF = np.stack(list(normal(xx, yy, h, p)[np.newaxis, ...] for p in points), axis=0)
    return dataF

def assemble_tensor(F, data_type, points):
    n = F.shape[-1]
    x = np.linspace(-1, 1, n)
    y = np.linspace(-1, 1, n)
    xx, yy = np.meshgrid(x, y)

    match data_type:
        case 'f':
            tensor = torch.from_numpy(F).float()
        
        case 'r':
            Rs = np.stack(list(_r(xx, yy, p)[np.newaxis, ...] for p in points), axis=0)
            tensor = torch.from_numpy(Rs).float()

        case 'rf':
            Rs = np.stack(list(_r(xx, yy, p)[np.newaxis, ...] for p in points), axis=0)
            tensor = torch.from_numpy(np.concatenate((Rs, F), axis=1))
        
        case 'xyf':
            tensor = np.stack(list(np.vstack([xx[np.newaxis, ...], yy[np.newaxis, ...], f]) for f in F), axis=0)
            tensor = torch.from_numpy(tensor)
    return tensor